# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment1/question3/interpolations/cubic_spline.py
import typing

import sympy

from question3.interpolation_method import InterpolationMethod1d


class CubicSplineInterpolation(InterpolationMethod1d):
    def interpolate(
            self,
            points: typing.Iterable[tuple[float, float]]) \
            -> typing.Callable[[float], float]:
        points = sorted(points, key=lambda p: p[0])
        if len(points) < 2:
            raise ValueError(f"Cannot do interpolation with {len(points)} point.")

        symbols: list[sympy.Symbol] = []
        equations = []

        points_iter = iter(points)
        previous_x, previous_y = next(points_iter)
        for i, (x, y) in enumerate(points_iter, 1):
            a, b, c, d = sympy.symbols(f"a{i} b{i} c{i} d{i}")
            # 线过点
            equations.append(
                a * x ** 3 + b * x ** 2 + c * x + d
                - y)
            equations.append(
                a * previous_x ** 3 + b * previous_x ** 2 + c * previous_x + d
                - previous_y)
            # 导连续
            if symbols:
                previous_c = symbols[-2]
                previous_b = symbols[-3]
                previous_a = symbols[-4]
                equations.append(
                    3 * a * previous_x ** 2 + 2 * b * previous_x + c
                    - 3 * previous_a * previous_x ** 2
                    - 2 * previous_b * previous_x
                    - previous_c)
                equations.append(
                    6 * a * previous_x + 2 * b
                    - 6 * previous_a * previous_x
                    - 2 * previous_b)
            symbols.append(a)
            symbols.append(b)
            symbols.append(c)
            symbols.append(d)
            previous_x = x
            previous_y = y

        # 首末三阶导为零
        equations.append(symbols[0])
        equations.append(symbols[-4])

        solution = sympy.linsolve(equations, symbols)

        result: dict[float, tuple[float, float, float, float]] = {}
        min_x = points[0][0]
        del points[0]

        solution_iter = iter(next(iter(solution)))
        for x, _ in points:
            a = float(next(solution_iter))
            b = float(next(solution_iter))
            c = float(next(solution_iter))
            d = float(next(solution_iter))
            result[x] = (a, b, c, d)
        pass

        def result_function(r_x: float):
            # TODO: 优化查找效率
            # 可以利用二分法等方式优化查找效率，而不是一个一个看。
            # 不过在插值点不是非常多的情况下，没有太大差别。
            # 且在 python 层进行这种优化可能有些尴尬（这样不能使用 C 底层提供的迭代器，还要进行额外的检查），
            # 所以这边暂时不作相关优化。
            if r_x < min_x:
                return float("nan")
            for k in result.keys():
                if r_x <= k:
                    r_a, r_b, r_c, r_d = result[k]
                    return r_a * r_x ** 3 + r_b * r_x ** 2 + r_c * r_x + r_d
            return float("nan")

        return result_function
